source("R_Imputation_Functions.R")
require("scater")
require("doMC")
require("e1071")
source("./Gmm/odgmm.R")
source("./Gmm/gmm_impute_core_functions.R")

set.seed(20194)
nrep=10;
my_seeds <- round(runif(nrep*2)*10000)
my_seeds <- unique(my_seeds)
my_seeds <- my_seeds[1:nrep]

# args <- commandArgs(trailingOnly=TRUE)

# method=args[1]


get_matrix <- function(sce, comp=2, seed=42, n.cores=10, verbose=FALSE) {
    
    mu0 = 0
    sigma0 = 0.1
    set.seed(seed)
    comp = comp
    X_tab = as.matrix(t(assays(sce)[["counts"]]))
    
	registerDoMC(n.cores)
	stime <- system.time({
	    res <- foreach(i=1:dim(X_tab)[1]) %dopar% {
	        X = as.numeric(X_tab[i,])   #X_tab is the matrix need to be fitted by GMM
	        
	        res = tryCatch(
	            odgmm(X,comp), 
	            error=function(e) {
	                if(verbose) print(e)
	                return(NA)
	            }
	        )
	        
	        if (length(res) == 0 || sum(is.na(unlist(res))) != 0) {
	            print(res)
	            return(NULL)
	        } else {
	            return(proc_res(res)) 
	        }
	    }
	    
	})[3]
	
	print(stime)
	
	X = as.data.frame(matrix(NA, nrow = dim(X_tab)[2], ncol = dim(X_tab)[1]))
	
	colnames(X) = rownames(X_tab)
	rownames(X) = colnames(X_tab)
	
	xvals = list()
	
	for (i in 1:length(res)) {
	    temp_res = res[[i]]
	    if(!is.null(temp_res)) {
	        X[, i] = temp_res[[1]]
	        xvals[[i]] = temp_res[[2]]
	    } 
	}
	
	X = na.omit(X)

	return(list(
		X=X,
		X_tab=X_tab,
		xvals=xvals
	))
}


perform_GMM_impute <- function(data, rate=0.8) {

    ## Second step:
    ## GMM output into imputing
	X = data[["X"]]
	X_tab = data[["X_tab"]]

	n_gene = dim(X)[1]
    n_cell = dim(X)[2]

	xvals = data[["xvals"]]

    keep_rate = apply(X>0,2,sum)/n_cell
    gene_inds = which(keep_rate > rate)

    print(summary(keep_rate))

    # perform imputation
    stime <- system.time({
        XX = X[,gene_inds]
        XX_vals = xvals[gene_inds]
        res_after_impute = gmm_impute(XX,XX_vals)
        X_impute = res_after_impute$impute_label_mat
        X_val_impute = res_after_impute$impute_val_mat # NA means values that does not need to be imputed
        orig_data = t(X_tab[gene_inds,])
        orig_data_verse = t(X_tab[-gene_inds,])
        valid_inds = is.na(X_val_impute)
        X_val_impute[valid_inds] = orig_data[valid_inds]
        GMM_imputed_X_tab = cbind(orig_data_verse,X_val_impute)
    })[3]
    
    print(paste0("Imputation takes: ", stime))

    return(t(GMM_imputed_X_tab))
}


perform_impute_test_test <- function(method, n.cores=10) {
	# wrapper functions
	if (method == "scImpute") {
		# my_seeds <- my_seeds[1:3]
		param_name <- "Dropout Threshold"
		param_range=rev(c(0, 0.2, 0.4, 0.6, 0.8, 1))
		imputation_fxn <- scImpute_wrapper
	} else if (method == "DrImpute") {
		param_name <- "Zeros Remaining"
		param_range=rev(c(0, 0.05, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.75, 1))
		imputation_fxn <- DrImpute_wrapper
	} else if (method == "SAVER") {
		# my_seeds <- my_seeds[1:3]
		param_name <- "Percent of Genes"
		param_range=c(-1, 0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1)
		imputation_fxn <- SAVER_wrapper
	} else if (method == "MAGIC") {
		param_name <- "K neighbours"
		param_range=c(-1, 5, 10, 15, 30, 45, 60, 100)
		imputation_fxn <- MAGIC_wrapper
	} else if (method == "knn") {
		param_name <- "K neighbours"
		param_range=c(-1, 5, 10, 15, 30, 45, 60, 100)
		imputation_fxn <- knn_wrapper
	} else if (method == "Gmm") {
		param_name <- "Gmm"
		param_range = rev(c(0.5, 0.52, 0.54, 0.56, 0.58))
		imputation_fxn <- GMM_impute_wrapper 
	} else if (method == "dca") {
		param_name <- "dca"
		param_range = c(8, 16, 24, 32, 40, 48, 56, 64)
	}

	AllTP <- vector()
	AllFP <- vector()
	AllTN <- vector()
	AllFN <- vector()

	print(param_name)
	print(param_range)

	for (s in my_seeds) {
		sims <- readRDS(paste("vsMethodParams_seed_",s,".rds", sep=""))

		require("Hmisc")
		TP <- vector()
		FP <- vector()
		TN <- vector()
		FN <- vector()

		if (method == "Gmm") {
			sims = get_matrix(sims, n.cores=n.cores)
			gc()
		}

		for (i in 1:length(param_range)) {
			print(paste("Seed: ", s, "; param: ", param_range[i], sep=""))
			old <- Sys.time() # get start time
			val <- param_range[i]

			if (method == "Gmm") {
				mat = perform_GMM_impute(sims, rate=param_range[i]) 
			} else if (method == "dca") {
				mat <- sims@assays[[paste("dca", val, sep = "_")]]
			} else if (val == -1) {
				mat <- sims@assays[["counts"]]
			} else if (!is.null(n.cores)) {
				mat <- imputation_fxn(sims, param=val, n.cores=n.cores)
			} else {
				mat <- imputation_fxn(sims, param=val)
			}
			cor_out <- Hmisc::rcorr(t(mat), type="spearman")
			thresh <- 0.05 / ( prod(dim(cor_out$P))/2 - length(diag(cor_out$P)) )
			de <- rowData(sims)$g_up | rowData(sims)$g_down
			TP <- c(TP, sum(cor_out$P[de, de] < thresh, na.rm=T))
			FP <- c(FP, sum(cor_out$P < thresh, na.rm=T) - TP[i])
			FN <- c(FN, sum(cor_out$P[de, de] >= thresh, na.rm=T))
			TN <- c(TN, sum(cor_out$P >= thresh, na.rm=T) - FN[i])
			new <- Sys.time() - old # calculate difference

			print(new) # print in nice format
			gc()
		}

		AllTP <- cbind(AllTP, TP)
		AllFP <- cbind(AllFP, FP)
		AllTN <- cbind(AllTN, TN)
		AllFN <- cbind(AllFN, FN)

	}

	FPR <- AllFP/(AllFP+AllTN)

	out <- list(TP=AllTP, FP=AllFP, TN=AllTN, FN=AllFN, params=param_range, param_name=param_name)
	saveRDS(out, file=paste(method,"vs_param.rds", sep="_"))
}


perform_imputetest("scImpute", n.cores=10)
perform_impute_test("DrImpute", n.cores=NULL)
perform_impute_test("SAVER", n.cores=10)
perform_impute_test("MAGIC", n.cores=NULL)
perform_impute_test("knn", n.cores=NULL)
perform_impute_test("Gmm", n.cores=10)

perform_impute_test("dca")

